import argparse
import multiprocessing
import time
import ray
import torch
import zmq
import tensordict

# Ensure this runs in the repository root directory, otherwise sys.path.append might be needed
from transfer_queue.storage.simple_backend import SimpleStorageUnit
from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType


class StorageClient:
    """Independent test client that interacts directly with the frontend ROUTER of SimpleStorageUnit"""

    def __init__(self, address):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.DEALER)
        self.socket.setsockopt(zmq.RCVTIMEO, 20000)  # Timeout set to 20s to prevent timeouts under heavy concurrency
        self.socket.connect(address)

    def send_put(self, client_id, local_indexes, field_data):
        msg = ZMQMessage.create(
            request_type=ZMQRequestType.PUT_DATA,
            sender_id=f"bench_client_{client_id}",
            body={"local_indexes": local_indexes, "data": field_data},
        )
        self.socket.send_multipart(msg.serialize())
        return ZMQMessage.deserialize(self.socket.recv_multipart())

    def close(self):
        self.socket.close()
        self.context.term()


def client_worker(worker_id, address, num_requests, batch_size):
    """Worker process task: Continuously bombard the Storage Unit with PUT requests"""
    client = StorageClient(address)
    start_time = time.time()

    # Construct Dummy Tensor data to simulate actual memory and serialization overhead
    # As noted in the PR description, serialization and memory I/O are the bottlenecks blocking the main loop
    field_data = {
        "dummy_tensor": [torch.randn(256, 256) for _ in range(batch_size)]
    }

    for i in range(num_requests):
        local_indexes = list(range(i * batch_size, (i + 1) * batch_size))
        client.send_put(worker_id, local_indexes, field_data)

    elapsed = time.time() - start_time
    client.close()

    print(f"[Worker {worker_id}] Completed {num_requests} write requests, took {elapsed:.3f} seconds "
          f"(QPS: {num_requests / elapsed:.2f} req/s)")


def main(num_clients, storage_threads, requests_per_client):
    # Initialize Ray and global settings
    ray.init(ignore_reinit_error=True)
    tensordict.set_list_to_stack(True).set()

    try:
        print(f"🚀 Launching SimpleStorageUnit, internal worker threads (num_worker_threads): {storage_threads} ...")

        # Launch the backend Actor. PR 37 exposes the num_worker_threads parameter
        storage_actor = SimpleStorageUnit.options(
            max_concurrency=50, num_cpus=2
        ).remote(
            storage_unit_size=1000000,
            num_worker_threads=storage_threads # comment this line for old version comparison
        )

        zmq_info = ray.get(storage_actor.get_zmq_server_info.remote())
        put_get_address = zmq_info.to_addr("put_get_socket")
        print(f"✅ Storage unit ready, ZMQ Address: {put_get_address}")

        # Wait for zmq.proxy and all worker threads to bind to the inproc port
        time.sleep(2)

        print(f"🔥 Spawning {num_clients} independent concurrent write processes...")
        processes = []
        batch_size = 256

        start_time = time.time()

        # 1. Create and start multiple processes
        for i in range(num_clients):
            p = multiprocessing.Process(
                target=client_worker,
                args=(i, put_get_address, requests_per_client, batch_size)
            )
            p.start()
            processes.append(p)

        # 2. Wait for all concurrent processes to complete
        for p in processes:
            p.join()

        total_time = time.time() - start_time
        total_requests = num_clients * requests_per_client

        print("\n" + "=" * 50)
        print(f" 📊 Benchmark Results")
        print("=" * 50)
        print(f" SimpleStorageUnit internal threads : {storage_threads}")
        print(f" External concurrent clients        : {num_clients}")
        print(f" Total processed requests (Batches) : {total_requests} (Batch Size: {batch_size})")
        print(f" Total benchmark duration           : {total_time:.3f} seconds")
        print(f" 🚀 Overall Throughput              : {total_requests / total_time:.2f} req/s")
        print("=" * 50 + "\n")

    finally:
        # Resource cleanup
        if 'storage_actor' in locals():
            ray.kill(storage_actor)
        ray.shutdown()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="PR #37 Performance Benchmark")
    parser.add_argument("--clients", type=int, default=8, help="Number of concurrent client processes")
    parser.add_argument("--threads", type=int, default=4, help="Number of processing threads in SimpleStorageUnit")
    parser.add_argument("--requests", type=int, default=300, help="Number of requests sent per client")

    args = parser.parse_args()
    main(args.clients, args.threads, args.requests)